use heapless::spsc::Queue;
use rayon::prelude::*;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

// 定义消息类型（用户可以自定义）
#[derive(Debug, Clone)]
pub enum Message {
    Start,
    Data(HashMap<String, String>),
    Stop,
}

// 定义服务端配置
pub struct ServerConfig {
    pub num_threads: usize,                                // 线程数
    pub handlers: Vec<Box<dyn Fn(Message) + Send + Sync>>, // 消息处理函数列表
}

impl ServerConfig {
    pub fn new(num_threads: usize) -> Self {
        Self {
            num_threads,
            handlers: Vec::new(),
        }
    }

    // 添加消息处理函数
    pub fn add_handler<F>(&mut self, handler: F)
    where
        F: Fn(Message) + Send + Sync + 'static,
    {
        self.handlers.push(Box::new(handler));
    }
}

// 定义服务端
pub struct Server {
    queue: Arc<Mutex<Queue<Message, 1024>>>,
    config: ServerConfig,
    shutdown_flag: Arc<Mutex<bool>>, // 关闭标志位
}

impl Server {
    pub fn new(queue: Arc<Mutex<Queue<Message, 1024>>>, config: ServerConfig) -> Self {
        Self {
            queue,
            config,
            shutdown_flag: Arc::new(Mutex::new(false)), // 初始化关闭标志位
        }
    }

    pub fn start(&self) {
        println!(
            "Server started with {} threads, waiting for messages...",
            self.config.num_threads
        );

        // 克隆关闭标志位
        let shutdown_flag = Arc::clone(&self.shutdown_flag);

        // 使用 rayon 线程池处理消息
        (0..self.config.num_threads).into_par_iter().for_each(|_| {
            loop {
                // 检查关闭标志位
                if *shutdown_flag.lock().unwrap() {
                    println!("Server: Shutdown flag detected, stopping...");
                    break;
                }

                // 从队列中读取消息
                let message = {
                    let mut queue = self.queue.lock().unwrap();
                    queue.dequeue()
                };

                if let Some(message) = message {
                    // 如果收到 Stop 消息，设置关闭标志位
                    if let Message::Stop = message {
                        *shutdown_flag.lock().unwrap() = true;
                        break;
                    }
                    // 调用所有注册的处理函数
                    for handler in &self.config.handlers {
                        handler(message.clone());
                    }
                } else {
                    // 如果队列为空，休眠一段时间
                    thread::sleep(Duration::from_millis(100));
                }
            }
        });

        println!("Server stopped.");
    }

    // 关闭服务端的方法
    pub fn shutdown(&self) {
        println!("Server: Shutting down...");
        *self.shutdown_flag.lock().unwrap() = true;
    }
}

// 定义客户端
pub struct Client {
    queue: Arc<Mutex<Queue<Message, 1024>>>,
}

impl Client {
    pub fn new(queue: Arc<Mutex<Queue<Message, 1024>>>) -> Self {
        Self { queue }
    }

    // 发送消息的方法
    pub fn send(&self, message: Message) {
        let mut queue = self.queue.lock().unwrap();
        if queue.enqueue(message).is_err() {
            println!("Client: Failed to send message, queue is full!");
        }
    }
}
